from alpha_vantage import fundamentaldata
import argparse
import pandas as pd
import random
import string
import tqdm
from typing import Any, Dict, List, Optional

_NASDAQ_FILE_NAME = 'nasdaq.txt'
_SP500_FILE_NAME = 'sp500.txt'
_TEST_FILE_NAME = 'test.txt'

_COLUMNS = ['Symbol', 'GrossMarginLatest', 'ProfitMarginLatest', 'ReturnOnEquityTTM',
            'ReturnOnAssetsTTM', 'DebtToAssetRatio', 'CurrentRatio', 'QuickRatio']

# Gross Margin、Net Margin、ROE、ROA
# Debt to Assets Ratio、Current Ratio、Quick Ratio

def parse_symbols():
    symbols = []
    with open('./sp500_table.txt', 'rt') as f:
        lines = f.readlines()
        for line in lines:
            symbols.append(line.split()[1])
    symbols.sort()
    with open('./sp500.txt', 'wt') as f:
        for symbol in symbols:
            f.write(f'{symbol}\n')


def generate_api_key(length: int = 16) -> str:
    characters = string.ascii_letters + string.digits
    random_string = ''.join(random.choice(characters) for i in range(length))
    return random_string


def process_data(fd: fundamentaldata.FundamentalData, file_name: str) -> pd.DataFrame:
    with open(file_name) as f:
        symbols = f.read().split()
    company_overview_list = []
    balance_sheet_list = []
    income_list = []
    num_iterations = len(symbols)

    for i in tqdm.tqdm(range(num_iterations), desc='Processing'):
        symbol = symbols[i].replace('.', '-')
        try:
            company_overview_result_temp = fd.get_company_overview(symbol)[0]
            balance_sheet_result_temp = pd.DataFrame(fd.get_balance_sheet_quarterly(symbol)[0].iloc[0]).T
            income_result_temp = pd.DataFrame(fd.get_income_statement_quarterly(symbol)[0].iloc[0]).T
        except Exception as e:
            print(f'Failed to retrieve data for {symbol}, error: {e}')
            continue
        company_overview_list.append(company_overview_result_temp)
        balance_sheet_result_temp['Symbol'] = symbol
        balance_sheet_list.append(balance_sheet_result_temp)
        income_result_temp['Symbol'] = symbol
        income_list.append(income_result_temp)

    company_overview_result = pd.concat(company_overview_list)

    balance_sheet_result = pd.concat(balance_sheet_list)
    balance_sheet_result['DebtToAssetRatio'] = (
        balance_sheet_result['totalLiabilities'].astype('float') / 
        balance_sheet_result['totalAssets'].astype('float'))
    total_current_liability = balance_sheet_result['totalCurrentLiabilities'].astype('float')
    total_current_assets = balance_sheet_result['totalCurrentAssets'].astype('float')
    inventory = pd.to_numeric(balance_sheet_result['inventory'], errors='coerce').fillna(0)
    balance_sheet_result['CurrentRatio'] = total_current_assets / total_current_liability
    balance_sheet_result['QuickRatio'] = (total_current_assets - inventory) / total_current_liability

    income_result = pd.concat(income_list)
    total_revenue = income_result['totalRevenue'].astype('float')
    gross_profit = income_result['grossProfit'].astype('float')
    net_income = income_result['netIncome'].astype('float')
    income_result['GrossMarginLatest'] = gross_profit / total_revenue
    income_result['ProfitMarginLatest'] = net_income / total_revenue
    tmp_result = pd.merge(company_overview_result, balance_sheet_result, on='Symbol')
    return pd.merge(tmp_result, income_result, on='Symbol')


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('-k', '--api_key', default='3LZJQ5CUCPSXJLLD')
    parser.add_argument('-f', '--file_name', default=_TEST_FILE_NAME)
    args = parser.parse_args()

    pd.options.mode.copy_on_write = True
    fd = fundamentaldata.FundamentalData(key=args.api_key, output_format='pandas')
    
    result = process_data(fd, args.file_name)
    df = result[_COLUMNS]
    df.rename(columns={'Symbol': '股票', 'GrossMarginLatest': '毛利率', 'ProfitMarginLatest': '净利率',
                       'ReturnOnEquityTTM': 'ROETTM', 'ReturnOnAssetsTTM': 'ROATTM',
                       'DebtToAssetRatio': '资产负债率', 'CurrentRatio': '流动比率',
                       'QuickRatio': '速动比率'}, inplace=True)
    df.to_csv('result.csv', index=False, float_format='%.2f')
    result.rename(columns={'Symbol': '股票', 'GrossMarginLatest': '毛利率', 'ProfitMarginLatest': '净利率',
                           'ReturnOnEquityTTM': 'ROETTM', 'ReturnOnAssetsTTM': 'ROATTM',
                           'DebtToAssetRatio': '资产负债率', 'CurrentRatio': '流动比率',
                           'QuickRatio': '速动比率'}, inplace=True)
    result.to_csv('all_results.csv', index=False, float_format='%.2f')


if __name__ == '__main__':
    main()